package msu.ml.data;

import msu.ml.core.*;
import weka.core.*;
import weka.classifiers.*;
import java.util.*;
import java.io.*;
import javax.imageio.*;
import java.awt.image.*;
import java.awt.*;


public class ImagePostProcessor extends DataPostProcessorAdapter
{
	public static final int IMAGE_DIM = 1500;

	private Color [] bioColor = {new Color(0xAFEEEE), new Color(0x00FFFF), 
		new Color(0x40E0D0), new Color(0x48D1CC), new Color(0x00CED1), new Color(0x00CED1)};
	
	private Color [] nonbioColor = {new Color(0xFFD39B), new Color(0xFFA54F),
		new Color(0xEE9A49), new Color(0xCD853F), new Color(0x8B5A2B), new Color(0x8B5A2B)};

	public ImagePostProcessor()
	{
	}

	public void beginExperiment(Experiment e)
	{
		super.beginExperiment(e);
	}

	public void beginInstances(String name) { }
	public void endInstances() { }
	public void process(Instance data, double[] distribution) { }

	public void endExperiment()
	{
		Experiment exp = getExperiment();
		Classifier c = exp.getClassifier();
		IDataProvider p = exp.getDataProvider();

		if(p instanceof IDbDataProvider)
		{
			IDbDataProvider dbp = 
				(IDbDataProvider)p;

			System.out.println("Creating Images...");

			HashSet<String> features = exp.getFeatures();
			boolean az = !features.contains("1");
			boolean ra = !features.contains("2");
			features.add("1");
			features.add("2");

			AttributeRemovalPreProcessor filter = new
				AttributeRemovalPreProcessor();
			filter.setInvertSelection(true);
			String featureString = features.toString();
			filter.setAttributeIndices(featureString.substring(1,
						featureString.length() - 1));

			ArrayList<IDataPreProcessor> filters = new
				ArrayList<IDataPreProcessor>();
			filters.add(filter);
			DatabaseCache cache = dbp.getCache();

			File imageDir = new File(exp.getName() + "_images");
			imageDir.mkdir();

			for(String file : dbp.getFileNames())
			{
				System.out.println(file + ".png");
				try
				{
					BufferedImage image = new BufferedImage(IMAGE_DIM, IMAGE_DIM, BufferedImage.TYPE_INT_ARGB);
					Graphics2D g = image.createGraphics();
					g.setColor(Color.white);
					g.fillRect(0,0,IMAGE_DIM,IMAGE_DIM);

					NxInstances richData = cache.retrieve(file, filters);
					NxInstances data = new NxInstances(richData);
					if(ra) data.deleteAttributeAt(1);
					if(az) data.deleteAttributeAt(0);

					for(int i = 0; i < data.numInstances(); i++)
					{
						double [] dist = c.distributionForInstance(data.instance(i));
						Instance inst = richData.instance(i);

						double r = inst.value(1) / 1000.0;
						double t = inst.value(0);

						int [] x = {cX(r, t), cX(r + 1, t), cX(r + 1, t + 1), cX(r, t + 1)};
						int [] y = {cY(r, t), cY(r + 1, t), cY(r + 1, t + 1), cY(r, t + 1)};

						g.setColor((dist[1] > dist[0]) ? bioColor[(int)(dist[1] * 10.0) - 5] 
								: nonbioColor[(int)(dist[0] * 10.0) - 5]);

						g.fillPolygon(x, y, 4);
					}

					File imageFile = new File(imageDir.getName() + File.separatorChar + file+".png");
					ImageIO.write(image, "png", imageFile);
				}
				catch(Exception e)
				{
					System.out.println(e);
					e.printStackTrace();
				}
			}
		}
	}

	private int cX(double r, double theta)
	{
		return (IMAGE_DIM / 2) - (int)(r * 5.0 * Math.cos(theta * Math.PI / 180.0));
	}

	private int cY(double r, double theta)
	{
		return (int)(r * 5.0 * Math.sin(theta * Math.PI / 180.0)) + (IMAGE_DIM / 2);
	}
}
